import os  
import numpy as np  
  
def find_centroid(points):  
    """计算点云的质心"""  
    centroid = np.mean(points[:, :3], axis=0)  
    return centroid  
  
def normalize_to_farthest_point(points, centroid):  
    """归一化点云，使得质心到最远点的距离为1"""  
    # 计算每个点到质心的距离  
    distances = np.linalg.norm(points[:, :3] - centroid, axis=1)  
    # 找到最远距离  
    max_distance = np.max(distances)  
    # 归一化点云坐标  
    normalized_points = np.hstack((  
        (points[:, :3] - centroid) / max_distance,  
        points[:, 3:]  # RGB值保持不变  
    ))  
    return normalized_points, max_distance  
  
def read_txt_file(filename):  
    """读取txt文件中的点云数据，包括xyz和rgb"""  
    points = []  
    with open(filename, 'r') as file:  
        for line in file:  
            values = line.split()  
            xyz = np.array([float(v) for v in values[:3]])  # x, y, z坐标  
            rgb = np.array([float(v) for v in values[3:6]])  # r, g, b值  
            point = np.hstack((xyz, rgb))  
            points.append(point)  
    return np.array(points)  
  
def write_txt_file(filename, points):  
    """将点云数据（包括归一化后的xyz和原始的rgb）写入txt文件"""  
    with open(filename, 'w') as file:  
        for point in points:  
            file.write(' '.join(map(str, point)) + '\n')  
  
# 输入文件夹路径，包含点云文件  
input_folder = 'point_cloud_models/sleeve'
# 输出文件夹路径，用于存储归一化后的点云文件  
output_folder = 'point_cloud_models/normalization'
  
# 确保输出文件夹存在  
if not os.path.exists(output_folder):  
    os.makedirs(output_folder)  
  
# 遍历输入文件夹中的所有txt文件  
for filename in os.listdir(input_folder):  
    if filename.endswith('.txt'):  
        # 构建输入和输出文件的完整路径  
        input_file_path = os.path.join(input_folder, filename)  
        output_file_path = os.path.join(output_folder, filename)  
          
        # 读取点云数据  
        points = read_txt_file(input_file_path)  
          
        # 找到质心并归一化点云  
        normalized_points, max_distance = normalize_to_farthest_point(points, find_centroid(points))  
          
        # 将归一化后的点云数据写入输出文件  
        write_txt_file(output_file_path, normalized_points)  
          
        print(f"Normalized {filename} with max distance {max_distance} has been written to {output_file_path}")  
  
print("All point clouds have been normalized and saved.")